import os
import pickle

import albumentations as A
import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm

from constants import FACE_FORENSICS
from training.datasets.face_forensics_all_dataset import MANIPULATION_TYPES
from training.datasets.transform_v2 import IsotropicResizeV2
from image_filtering import image_gradient_enhance

transform = A.Compose([
    IsotropicResizeV2(min_side=256),
    A.CenterCrop(height=256, width=256),
])


def mse(imageA, imageB):
    # the 'Mean Squared Error' between the two images is the
    # sum of the squared difference between the two images;
    # NOTE: the two images must have the same dimension
    err = np.sum((imageA.astype("float") / 255. - imageB.astype("float") / 255.) ** 2)
    err /= float(imageA.shape[0] * imageA.shape[1])

    # return the MSE, the lower the error, the more "similar"
    # the two images are
    return err


def main():
    os.makedirs('analysis', exist_ok=True)
    root_dir = '/home/xinlin/data2/FaceForensics++'
    version = 'c40'
    original_path = os.path.join(root_dir, 'original_sequences', 'youtube', version)

    for fake_type in MANIPULATION_TYPES:
        csv = f'../data/{FACE_FORENSICS}/c40/data_{FACE_FORENSICS}_{fake_type}_train_fake.csv'
        df = pd.read_csv(csv, converters={'video': lambda x: str(x), 'original': lambda x: str(x)})
        # df = df[df['label'] == 1]
        # df_first = df.groupby('video').first()
        df_sample = df.sample(n=720 * 5, random_state=111)
        df_sample['fake_type'] = fake_type
        manipulated_path = os.path.join(root_dir, 'manipulated_sequences', fake_type, version)

        mse_scores = []
        for index, row in tqdm(df_sample.iterrows()):
            # video = str(index)
            video, img_file, label, ori_video, frame, fake_type = row.values

            fake_path = os.path.join(manipulated_path, 'crops', video, img_file)
            real_path = os.path.join(original_path, 'crops', ori_video, img_file)
            fake_image = cv2.imread(fake_path, cv2.IMREAD_COLOR)
            real_image = cv2.imread(real_path, cv2.IMREAD_COLOR)
            fake_image = transform(image=fake_image)['image']
            real_image = transform(image=real_image)['image']
            mse_score = mse(fake_image, real_image)

            fake_image_sobel = image_gradient_enhance(fake_image)
            real_image_sobel = image_gradient_enhance(real_image)
            mse_score_sobel = mse(fake_image_sobel, real_image_sobel)
            # print(mse_score, mse_score_sobel)
            mse_scores.append((mse_score, mse_score_sobel))

        with open(f'./analysis/analysis_{fake_type}.pickle', "wb") as f:
            pickle.dump(mse_scores, f)


if __name__ == '__main__':
    main()
